'''
程序名：File_AES.py：
功能：Python 实现 文件AES加密
编程人：Rzz

加密文件头架构：后缀名：“.FAC”（此文件头也被加密！）
文件特征符（12 Byte：b'File_AES\x55\xaa\x33\xcc'）
文件名长度（4 Byte: b(int)）
文件名（n Byte, encode'utf-8'）
哈希值（32 Byte: 明文文件SHA256值）
保密正文（xxxx Byte: 与明文文件同长度）

加密配置文件名："File_AES.ini"。架构如下：
文件特征符（12 Byte：b'File_AES\x55\xaa\x33\xcc'）
长密码长度（4 Byte: b(int)）
随机数凑成：64 Byte
'''
#!/usr/bin/env python3
# coding:utf-8

import tkinter # GUI 图形用户界面
from tkinter import filedialog  # 文件操作库
from tkinter import messagebox  # 单独引入消息窗口库
from tkinter import simpledialog  # 单独引入信息输入窗口库

from Cryptodome.Cipher import AES
from Cryptodome import Random
from Cryptodome.Hash import SHA256

import time # 用于进度条显示时间
import os # 用于对文件名和路径的处理


class GuiWin: # 游戏界面类，含控制流程
    fhead_index = b'File_AES\x55\xaa\x33\xcc' # 12位文件标识符，用于检查解密是否正确
    fhindex_len = len(fhead_index) # 文件头引导符长度:12
    fnlen_len = 4 # 预留固定4字节，存放文件名的长度
    fhead_len = fhindex_len + fnlen_len # 整体长度:16
    hash_len = 32 # 文件hash值长度

    _safeflag = False # 安全模式标志，初始为非！
    fin_name = "" # 设置文件名初始值为空，避免程序错误提示
#    real_key =  # 密码运算实际key，程序中直接赋值
#    real_nonce1 = # 密码运算实际nonce1，程序中直接赋值
#    real_nonce2 = # 密码运算实际nonce2，程序中直接赋值
    
    def __init__(self): # 显示界面类初始化：显示窗体
        self.window = tkinter.Tk() # 创设主窗口
        self.window.title("文件AES加密")

        sw = self.window.winfo_screenwidth() # 获取屏幕宽度
        sh = self.window.winfo_screenheight() # 获取屏幕高度
        ww, wh = 300, 280 # 设置窗口宽度、高度
        x, y = (sw-ww)/2, (sh-wh)/2
        self.window.geometry("%dx%d+%d+%d" %(ww,wh,x,y)) # 设置窗口在屏幕中间


        # 文件操作区
        frm_f_r, frm_f_c = 0, 0 # 设置frame初始行列设置
        frm_f = tkinter.LabelFrame(self.window, width=300, height=90, text='文件操作').grid(row=frm_f_r, rowspan=4, column=frm_f_c, columnspan=8) # 创建文件操作框架区

        tkinter.Button(frm_f, text="输入文件",command = self.select_sfile).grid(row=frm_f_r+2)
        self.SourceLabel = tkinter.Label(frm_f, width=31, text="输入文件名", justify = "right", anchor = "w")
        self.SourceLabel.grid(row=frm_f_r+2, column=frm_f_c+1, columnspan=6)

#        tkinter.Button(frm_f, text="输出文件",command = self.select_dfile).grid(row=frm_f_r+3)
        tkinter.Label(frm_f, text="输出文件").grid(row=frm_f_r+3)
#        self.DistLabel = tkinter.Label(frm_f, width=31, text="输出文件名", anchor = "w")
        self.DistLabel = tkinter.Label(frm_f, width=31, text="", anchor = "w")
        self.DistLabel.grid(row=frm_f_r+3, column=frm_f_c+1, columnspan=6)

        # 密码输入区
        frm_k_r, frm_k_c = 4, 0 # 设置frame初始行列设置
        frm_k = tkinter.LabelFrame(self.window, width=300, height=120, text='密码输入').grid(row=frm_k_r, rowspan=4, column=frm_k_c, columnspan=8) # 创建文件操作框架区

#        tkinter.Label(frm_k, text="输入密码：").grid(row=frm_k_r+1, column=0)
        self.firstl = tkinter.Label(frm_k, text="输入密码：")
        self.firstl.grid(row=frm_k_r+1, column=0)
        self.key_v1 = tkinter.Variable() 
        self.key_v1.set("") # 设置文本框中的值,获取值函数：var.get()
#        tkinter.Entry(frm_k, textvariable=self.key_v1, show="*").grid(row=frm_k_r+1, column=1, columnspan=6, sticky='W E')
        self.tk_entry1 = tkinter.Entry(frm_k, textvariable=self.key_v1, show="*")
        self.tk_entry1.grid(row=frm_k_r+1, column=1, columnspan=6, sticky='W E')

        self.secondl = tkinter.Label(frm_k, text='再次输入：')
        self.secondl.grid(row=frm_k_r+2, column=0) # 必须分行，否者参数不能被修改！！！！！
        self.key_v2 = tkinter.Variable() 
        self.key_v2.set("") # 设置文本框中的值,获取值函数：var.get()
        self.tk_entry2 = tkinter.Entry(frm_k, textvariable=self.key_v2, show="*")
        self.tk_entry2.grid(row=frm_k_r+2, column=1, columnspan=6, sticky='W E') # 必须分行，否者参数不能被修改！！！！！

#        tkinter.Label(frm_k, text="处理方式").grid(row=frm_k_r+3, column=0) # 选择文件加密类型

        self.select_v = tkinter.IntVar()
        self.select_v.set(0) # 设置起始默认选项
        self._isen = True # 初始设定为加密
        self.rb1=tkinter.Radiobutton(frm_k, text="加密", variable=self.select_v, value=0, command = self.en_or_de)
        self.rb1.grid(row=frm_k_r+3, column=1) # 必须分行，否者参数不能被修改！！！！！
        self.rb2=tkinter.Radiobutton(frm_k, text="解密", fg='gray', variable=self.select_v, value=1, command = self.en_or_de)
        self.rb2.grid(row=frm_k_r+3, column=3) # 必须分行，否者参数不能被修改！！！！！


        # 执行区
        frm_r_r, frm_r_c = 8, 0 # 设置frame初始行列设置
        frm_r = tkinter.Frame(self.window, width=300, height=80).grid(row=frm_r_r, rowspan=4, column=frm_r_c, columnspan=8) # 创建文件操作框架区
        tkinter.Button(frm_r, text="运行",command = self.main_run).grid(row=frm_r_r+1, column=2)
        tkinter.Button(frm_r, text="简化模式",command = self.safe_mode).grid(row=frm_r_r+1, column=4) # 启动"安全模式"按钮

        # 版权说明
        tkinter.Label(frm_r,text=" Designed by: Rzz ").grid(row=frm_r_r+2, column=3, columnspan=4, sticky='E')

    def select_sfile(self): # 选择输入文件
        self.fin_name = filedialog.askopenfilename(title='打开输入文件',  # 直接设置类属性，为大家公用
                                                    initialdir='C:\Python_test',
                                                    filetypes=[("XLSX",".xlsx"),
                                                               ("FAC",".fac"),
                                                               ("*.*",".*")])
        self.SourceLabel['text'] = self.fin_name
        self.DistLabel['text'] = '' # 清除输出文件名
        
    def select_dfile(self): # 确定输出文件，本程序未用上！！！
        self.fout_name = filedialog.asksaveasfilename(title='确定输出文件',  # 直接设置类属性，为大家公用
                                                    initialdir='C:\Python_test',
                                                    filetypes=[("TXT",".txt"),
                                                               ("FAC",".fac"),
                                                               ("ERR",".err")])
        self.DistLabel['text'] = self.fout_name 

    def en_or_de(self): # 加密、解密选择
        if self.select_v.get() == 0:
            self._isen = True # 确认为加密
            if self._safeflag == False:
                self.tk_entry2['state'] = 'normal'
                self.secondl['fg'] = 'black'
            self.rb1['fg'] = 'black'
            self.rb2['fg'] = 'gray'
        else:
            self._isen = False # 确认为解密
            self.tk_entry2['state'] = 'disabled'
            self.secondl['fg'] = 'gray'
            self.rb1['fg'] = 'gray'
            self.rb2['fg'] = 'black'

    def safe_mode(self): # 安全模式。在安全环境下，日常使用短密吗替换长密码，提升效率
#        self._key3v = "" # 初值长度为零
#            self.DistLabel['text'] = self._key3v
#        _fname = os.path.join(os.getcwd(), "File_AES.ini")
        _fname = "File_AES.ini" # 长密码配置文件名
        
        if (os.path.isfile(_fname) == True): # 判断"File_AES.ini"文件是否存在?
            self._safeflag = True # 设置安全模式为True

            _fin = open(_fname, 'rb') # 打开输入文件
            _fin.seek(0,2) # 指针转到文件尾
            _fin_len = _fin.tell() # 获取文件长度
            _fin.seek(0,0) # 指针回到文件头
            if (_fin_len == 0): # 若文件长度为零，则设定安全环境密码
                _fin.close() # 首先关闭文件
                _fout = open(_fname,'wb') # 打开输出文件
#                self._key3v = simpledialog.askstring('密码输入框', '请输入密码：', show='*').encode('utf-8') # 编码后用户输入密码
                while True: # 强制等待输入密码
                    self._key3v = simpledialog.askstring('密码输入框', '输入长密码：', show='*').encode('utf-8') # 编码后用户输入密码
                    if ((len(self._key3v) < 8) or (len(self._key3v) > 32)): 
                        messagebox.showwarning('密码长度<8 or >32', '请重新输入密码！')
                    else :
                        self._key4v = simpledialog.askstring('密码输入框', '再次输入长密码：', show='*').encode('utf-8') # 编码后用户输入密码
                        if (self._key4v == self._key3v):
                            break
                        else :
                            messagebox.showwarning('两次输入密码不同', '请重新输入密码！')

                while True: # 强制等待输入简化密码
                    self._key4v = simpledialog.askstring('密码输入框', '输入简化密码：', show='*').encode('utf-8') # 编码后用户输入密码
                    if (len(self._key4v) < 3): 
                        messagebox.showwarning('密码长度不够3位', '请重新输入密码！')
                    else:
                        break

                self.key_v1.set(self._key4v) # 传递短密码
                self.rebuild_key() # 重构Key及nonce
                AE.init(self.real_key, self.real_nonce1) # 初始化加密模块
#                b_str = AE.crypt(self._key3v)
                _fout_d = self.fhead_index + bytes('%4d'%len(self._key3v),encoding="utf-8") + self._key3v
                _fout_d += Random.get_random_bytes(64 - len(_fout_d)) # 长度随机添加成64位
                _fout.write(AE.crypt(_fout_d)) # 加密长密码文件
                _fout.close() # 关闭文件

#                self.DistLabel['text'] = self._key3v
            else: # 若长度不为零，则读取文件，获取执行密码
                self._key3v = _fin.read(_fin_len)
                while True: # 强制等待输入简化密码
#                    self._key4v = simpledialog.askstring('密码输入框', '输入简化密码：', show='*').encode('utf-8') # 编码后用户输入密码
                    _temp = simpledialog.askstring('密码输入框', '输入简化密码：', show='*')
                    if _temp: # 若输入非None
                        self._key4v = _temp.encode('utf-8') # 编码后用户输入密码
                    else:
                        self.key_v1.set("") # 清空密码
                        return
                    if (len(self._key4v) < 3): 
                        messagebox.showwarning('密码长度不够3位', '请重新输入密码！')
                    else:
                        break

                self.key_v1.set(self._key4v) # 传递短密码
                self.rebuild_key() # 重构Key及nonce
                AE.init(self.real_key, self.real_nonce1) # 初始化加密模块
                self._key3v = AE.crypt(self._key3v) # 解密长密码
                _fin.close() # 关闭文件
                if self._key3v[0:len(self.fhead_index)] == self.fhead_index: # 判断文件头是否解码正确，否则为密码错误！
                    self._key3v = self._key3v[len(self.fhead_index):] # 去除文件头
                    _len = int(self._key3v[:4]) # 获得密码长度
                    self._key3v = self._key3v[4:] # 去除密码长度
#                    self.DistLabel['text'] = self._key3v # for test!
                    self._key3v = self._key3v[:_len] # 去除随机数据
                else:
                    messagebox.showwarning('简化密码解码错误！', '请重新输入密码！')
                    self._safeflag = False # 设置安全模式为 False
                    self.key_v1.set("") # 清空密码
                    return

            self.tk_entry1['state'] = 'disabled'
            self.tk_entry2['state'] = 'disabled'
            self.firstl['fg'] = 'gray'
            self.secondl['fg'] = 'gray'
            self.key_v1.set(self._key3v) # 传递长密码
            self.key_v2.set(self._key3v) # 传递长密码
        else:
            if (messagebox.askyesno('环境确认', '请确认是否为安全环境？')): # 确认是否为安全环境？
                _fout = open(_fname,'wb') # 打开输出文件
                _fout.close() # 关闭文件
#            else:    
#                self.DistLabel['text'] = "非安全环境！"
      

    def __check(self): # 执行前检查
        if (self.fin_name == ''): # 没有设置文件名，直接返回错误
            messagebox.showwarning('文件名为空', '请选择输入文件！')
            return False
        if (self.key_v1.get() == ''): # 没有输入密码，直接返回错误
            messagebox.showwarning('未输入密码', '请输入密码！')
            return False
        if self._isen == True: # 若是加密，需要判断密码长度和验证对比
            if (len(self.key_v1.get()) < 8): # 密码长度不够，直接返回错误
                messagebox.showwarning('密码长度不够', '请重新输入密码！')
                return False
            if (self.key_v1.get() != self.key_v2.get()): # 没有输入密码，直接返回错误
                messagebox.showwarning('两次输入密码不同', '请重新输入密码！')
                return False
        return True

    def rebuild_key(self): # 重建秘钥
        if self._safeflag == False: # 非安全模式需要重新编码
            _d0 = self.key_v1.get().encode('utf-8') # 编码后用户输入密码
        else:
            _d0 = self.key_v1.get()
            
        _sha = SHA256.new() # 哈希算法初始化，用于重构key
        _sha.update(_d0 + AE.start_key)
        self.real_key = _sha.digest() # 获得哈希值，长度32byte

        _sha = SHA256.new() # 哈希算法初始化，用于重构nonce
        _sha.update(_d0 + AE.start_nonce)
        _d0 = _sha.digest() # 获得哈希值，长度32byte
        self.real_nonce1 = _d0[:12]
        self.real_nonce2 = _d0[20:]

    def __crypt_file(self, fin_len): # 执行前检查
        b_len = AE.b_len # 获取加密数据块长度
        _cnt = fin_len // b_len # 获得循环次数

        _sha = SHA256.new() # 哈希算法初始化
        PB = ProgressBar() # 实体化进度条
        time_start = time.time()

        for _i in range(_cnt+1): # 循环加密
            if _i == _cnt: # 尾数处理
                b_len == fin_len % b_len
            _din = self.fin.read(b_len)
            _dout = AE.crypt(_din)
            self.fout.write(_dout)
            time_end = time.time()
            PB.change(_i, _cnt, time_end - time_start)

            if self._isen == True: # 是加密么？
                _sha.update(_din)
            else:
                _sha.update(_dout)

        time.sleep(0.5) # 暂停
        PB.pb_destroy() # 关闭进度条
        return _sha.digest() # 返回文件hash值

    def __encrypt(self): # 加密文件
        _fn_sp = os.path.splitext(self.fin_name) # 去除后缀名
        self.fout_name = _fn_sp[0] + '.fac' # 合成带路径输出文件名

        self.fin_bn = os.path.basename(self.fin_name) # 不带路径的输入文件名
        bfin_bn = bytes(self.fin_bn, encoding = "utf8") # 输入文件名，编码为"utf8"格式
        _fn_sp = os.path.splitext(self.fin_bn) # 去除后缀名
        _fout_bn = _fn_sp[0] + '.fac' # 不带路径的输出文件名，用于界面提醒！

        if (os.path.isfile(self.fout_name) == True): # 判断输出文件是否存在
            if (messagebox.askyesno(_fout_bn + '文件存在！', '是否直接覆盖？') == False):
                return False
        self.DistLabel['text'] = self.fout_name 

        self.fin = open(self.fin_name, 'rb') # 打开输入文件
        self.fout = open(self.fout_name,'wb') # 打开输出文件

        self.fin.seek(0,2)
        fin_len = self.fin.tell() # 获取文件长度
        self.fin.seek(0,0)

        _len = len(bfin_bn) # 编码后输入文件名长度
        b_str = self.fhead_index + bytes(('%4d'%_len),encoding="utf-8") + bytes(bfin_bn)

        self.fout.write(bytes(self.fhead_len + _len + self.hash_len)) # 填写空字节，占位子

        AE.init(self.real_key, self.real_nonce2) # 初始化加密模块
        _hash = self.__crypt_file(fin_len) # 加密文件
        b_str += _hash

        self.fout.seek(0,0) # 调整指针，写文件头
        AE.init(self.real_key, self.real_nonce1) # 初始化加密模块
        b_str = AE.crypt(b_str)
        self.fout.write(b_str) # 加密文件头

        self.fin.close()
        self.fout.close()

    def __decrypt(self): # 解密文件
        self.fin = open(self.fin_name, 'rb') # 打开输入文件
        _data = self.fin.read(1024) # 预读1K数据，解密文件头

        AE.init(self.real_key, self.real_nonce1) # 初始化加密模块
        _data = AE.crypt(_data) # 解密文件头

        if _data[0:self.fhindex_len] != self.fhead_index: # 验证文件头特征符
            messagebox.showwarning('解码错误', '请重新输入密码！')
            self.fin.close()
            return False
            
        _len = int(_data[self.fhindex_len:self.fhead_len]) # 获得文件名长度
        self.fout_name = _data[self.fhead_len:self.fhead_len+_len].decode(encoding='utf-8')
        self.fout_name = os.path.join(os.path.dirname(self.fin_name),self.fout_name) 
        _fout_bn = os.path.basename(self.fout_name)

        if (os.path.isfile(self.fout_name) == True): # 判断输出文件是否存在
            if (messagebox.askyesno(_fout_bn + '文件存在！', '是否直接覆盖？') == False):
                self.fin.close()
                return False

        self.DistLabel['text'] = self.fout_name 
        self.fout = open(self.fout_name,'wb') # 打开输出文件

        self.fin.seek(0,2)
        fin_len = self.fin.tell() # 获取文件长度
        _len += self.fhead_len + self.hash_len # 获取文件头总长度
        fout_len = fin_len - _len # 去除文件头总长度
        self.fin.seek(_len,0)

        AE.init(self.real_key, self.real_nonce2) # 初始化加密模块
        _hash = self.__crypt_file(fout_len) # 加密文件

        self.fin.close()
        self.fout.close()

        _len -= self.hash_len # 获取hash值起始位置
        if _hash != _data[_len:_len+self.hash_len]: # 文件被篡改？
            if (messagebox.askyesno('加密文件被篡改！', '是否删除输出文件？') == True):
                os.remove(self.fout_name)

    def main_run(self): # 键盘输入判断，转方向事件程序处理
        self.DistLabel['text'] = ''
        if self.__check() == False: # 执行前检查未通过
            return

        self.rebuild_key() # 重构Key及nonce
        if self._isen == True:
            self.__encrypt()
        else:
            self.__decrypt()


class ProgressBar: # 进度条类, 单独列出，便于后续程序借用
    pb_w = 200 # 设置进度条宽度

    def __init__(self): # 类初始化：显示窗体
        self.top_win = tkinter.Toplevel() #bg = "grey")
        self.top_win.title('进度条')

        sw = self.top_win.winfo_screenwidth() # 获取屏幕宽度
        sh = self.top_win.winfo_screenheight() # 获取屏幕高度
        ww, wh = self.pb_w + 10, 60 # 设置窗口宽度、高度
        x, y = (sw-ww)/2, (sh-wh)/2
        self.top_win.geometry("%dx%d+%d+%d" %(ww,wh,x,y)) # 设置窗口在屏幕中间

        self.canvas = tkinter.Canvas(self.top_win,width = self.pb_w + 10,height = 30,bg = "white")
        self.canvas.grid(row = 1,column = 0)

        self.out_rec = self.canvas.create_rectangle(5,5,self.pb_w+5,25,outline = "blue",width = 1)
        self.fill_rec = self.canvas.create_rectangle(5,5,5,25,outline = "",width = 0,fill = "blue")
         
        self.x = tkinter.StringVar()
        tkinter.Label(self.top_win,textvariable = self.x).grid(row = 2,column = 0)

    def change(self,_now,_all,_time):
        _data = _now/_all
        self.canvas.coords(self.fill_rec, (5, 5, 5 + _data*self.pb_w, 25))
        if _now == _all:
            self.x.set("完成")
        elif _now == 0:
            self.x.set("开始...")
        else:
            self.x.set('已完成：'+str(round(_data*100,1)) + '%'+' 还需: %.2f s'%((_time/_data)*(1-_data)))
        self.top_win.update()

    def pb_destroy(self):
        self.top_win.destroy()


class Aes_Encrypt: # AES加密处理类
    start_key = b'The 16 bytes key' # 16 bytes初始key(16\24\32)
    start_nonce = b'12bytesnonce' # 8-15 bytes初始nonce
    b_len = 16 * 64 # 设置合适的加密数据块长度，最好为16的整数倍

    def init(self, _key, _nonce): # 使用key和nonce初始化AES对象, 使用MODE_CTR模式
        self.cipher = AES.new(_key, AES.MODE_CTR, nonce=_nonce) 

    def crypt(self, _data): # 实施加密/解密为同样计算路径
        return self.cipher.encrypt(_data) 

        
if __name__ == '__main__': # 如果是直接从本程序执行，则启动程序
    AE = Aes_Encrypt() # 实例化加密类
    GW = GuiWin() # 实例化GUI类 
    GW.window.mainloop() # 进入消息循环        
